#include "common.h"
#include "device/mmio.h"
#include "memory/memory.h"
#include "memory/l2cache.h"
#include "misc.h"
#include <stdlib.h>
#include <inttypes.h>

#define L2_NUM_SETS (L2_CACHE_SIZE / (L2_CACHE_BLOCK_SIZE * L2_CACHE_ASSOC))

typedef struct {
    bool valid;
    bool dirty;
    uint32_t tag;
    uint8_t data[L2_CACHE_BLOCK_SIZE];
} L2Line;

static L2Line l2_sets[L2_NUM_SETS][L2_CACHE_ASSOC];
static uint64_t l2_hits = 0, l2_misses = 0;

static inline uint32_t l2_off(hwaddr_t a) { return a & (L2_CACHE_BLOCK_SIZE - 1); }
static inline uint32_t l2_idx(hwaddr_t a) { return (a / L2_CACHE_BLOCK_SIZE) % L2_NUM_SETS; }
static inline uint32_t l2_tag(hwaddr_t a) { return (uint32_t)(a / L2_CACHE_BLOCK_SIZE) / L2_NUM_SETS; }

void l2_cache_init(void) {
    int i, j;
    for (i = 0; i < L2_NUM_SETS; i++) {
        for (j = 0; j < L2_CACHE_ASSOC; j++) {
            l2_sets[i][j].valid = false;
            l2_sets[i][j].dirty = false;
            l2_sets[i][j].tag = 0;
        }
    }
}

void l2_cache_reset_stats(void) { l2_hits = l2_misses = 0; }

static inline uint32_t l2_choose_victim(uint32_t set) {
    int j;
    for (j = 0; j < L2_CACHE_ASSOC; j++) if (!l2_sets[set][j].valid) return j;
    return rand() % L2_CACHE_ASSOC;
}

static int l2_probe_or_fill(uint32_t set, uint32_t tag, hwaddr_t block_base, int* hit) {
    int j;
    for (j = 0; j < L2_CACHE_ASSOC; j++) {
        if (l2_sets[set][j].valid && l2_sets[set][j].tag == tag) { *hit = 1; return j; }
    }
    *hit = 0;
    int v = l2_choose_victim(set);
    if (l2_sets[set][v].valid && l2_sets[set][v].dirty) {
        uint32_t old_tag = l2_sets[set][v].tag;
        hwaddr_t old_block_base = ((hwaddr_t)((old_tag * L2_NUM_SETS) + set)) * L2_CACHE_BLOCK_SIZE;
        int k;
        for (k = 0; k < L2_CACHE_BLOCK_SIZE; k += 4) {
            uint32_t w = (uint32_t)l2_sets[set][v].data[k] | ((uint32_t)l2_sets[set][v].data[k + 1] << 8)
                | ((uint32_t)l2_sets[set][v].data[k + 2] << 16) | ((uint32_t)l2_sets[set][v].data[k + 3] << 24);
            extern void dram_write(hwaddr_t, size_t, uint32_t);
            dram_write(old_block_base + k, 4, w);
        }
        l2_sets[set][v].dirty = false;
    }
    int k;
    extern uint32_t dram_read(hwaddr_t, size_t);
    for (k = 0; k < L2_CACHE_BLOCK_SIZE; k += 4) {
        uint32_t val = dram_read(block_base + k, 4);
        l2_sets[set][v].data[k + 0] = (uint8_t)(val & 0xff);
        l2_sets[set][v].data[k + 1] = (uint8_t)((val >> 8) & 0xff);
        l2_sets[set][v].data[k + 2] = (uint8_t)((val >> 16) & 0xff);
        l2_sets[set][v].data[k + 3] = (uint8_t)((val >> 24) & 0xff);
    }
    l2_sets[set][v].valid = true;
    l2_sets[set][v].dirty = false;
    l2_sets[set][v].tag = tag;
    return v;
}

int l2_cache_read_block(hwaddr_t block_base_addr, uint8_t* dst) {
    int map_NO = is_mmio(block_base_addr);
    if (map_NO != -1) {
        int k;
        for (k = 0; k < L2_CACHE_BLOCK_SIZE; k += 4) {
            uint32_t v = mmio_read(block_base_addr + k, 4, map_NO);
            dst[k + 0] = (uint8_t)(v & 0xff);
            dst[k + 1] = (uint8_t)((v >> 8) & 0xff);
            dst[k + 2] = (uint8_t)((v >> 16) & 0xff);
            dst[k + 3] = (uint8_t)((v >> 24) & 0xff);
        }
        return 0;
    }
    uint32_t set = l2_idx(block_base_addr);
    uint32_t tag = l2_tag(block_base_addr);
    int hit;
    int way = l2_probe_or_fill(set, tag, block_base_addr, &hit);
    if (hit) l2_hits++; else l2_misses++;
    memcpy(dst, l2_sets[set][way].data, L2_CACHE_BLOCK_SIZE);
    return hit;
}

void l2_cache_write(hwaddr_t addr, size_t len, uint32_t data) {
    Assert(len == 1 || len == 2 || len == 4, "len must be 1/2/4");
    int map_NO = is_mmio(addr);
    if (map_NO != -1) { mmio_write(addr, len, data, map_NO); return; }

    uint32_t off = l2_off(addr);
    if (off + len > L2_CACHE_BLOCK_SIZE) {
        size_t rem = len; hwaddr_t cur = addr; uint32_t shift = 0;
        while (rem > 0) {
            uint32_t off2 = l2_off(cur);
            size_t boundary = L2_CACHE_BLOCK_SIZE - off2;
            size_t chunk = rem; if (chunk > boundary) chunk = boundary; if (chunk >= 2) chunk = 2; else chunk = 1;
            uint32_t part = (data >> shift) & (~0u >> ((4 - chunk) << 3));
            l2_cache_write(cur, chunk, part);
            cur += chunk; shift += 8 * chunk; rem -= chunk;
        }
        return;
    }

    uint32_t set = l2_idx(addr);
    uint32_t tag = l2_tag(addr);
    hwaddr_t block_base = addr - off;
    int hit; int way = l2_probe_or_fill(set, tag, block_base, &hit);
    if (hit) l2_hits++; else l2_misses++;

    {
        uint8_t mask_local[4] = { 0 }; int i; for (i = 0; i < (int)len; i++) mask_local[i] = 1;
        memcpy_with_mask(l2_sets[set][way].data + off, &data, len, mask_local);
        l2_sets[set][way].dirty = true;
    }
}

void l2_cache_report_stats(void) {
    uint64_t acc = l2_hits + l2_misses;
    uint32_t rate = acc ? (uint32_t)(l2_hits * 10000 / acc) : 0;
    printf("[L2] size=%uB, block=%uB, ways=%u | accesses=%" PRIu64 ", hits=%" PRIu64 ", misses=%" PRIu64 ", hit_rate=%u.%02u%%\n",
        (unsigned)L2_CACHE_SIZE, (unsigned)L2_CACHE_BLOCK_SIZE, (unsigned)L2_CACHE_ASSOC,
        acc, l2_hits, l2_misses, rate / 100, rate % 100);
}
